import os
import json

import pytest

from test.test_utils import LOGGER, EnhancedJSONEncoder, is_huggingface_image

# Required to prevent circular dependency while importing
from test.test_utils import ecr as ecr_utils
from test.test_utils.security import (
    CVESeverity,
    ECREnhancedScanVulnerabilityList,
    generate_future_allowlist,
)


def get_object_after_serialization(input_object):
    """
    Serializes the object and returns its json equivalent
    """
    return json.loads(json.dumps(input_object, cls=EnhancedJSONEncoder))


# allowlist1.json represents the data exactly as it would look after being read into the ScanVulnerability Object from ecr_scan_result1.json.
# allowlist2.json represents the data exactly as it would look after being read into the ScanVulnerability Object from ecr_scan_result2.json.
# allowlist_to_read.json represents how the allowlist files would look on github repo.
# dump_of_processed_allowlist_to_read.json represents the data exactly as it would look after allowlist_to_read.json file is read into the ScanVulnerability Object.


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration("Check if ECR Scan data is read properly")
def test_read_from_ecr_scan():
    """
    This method tests that construct_allowlist_from_ecr_scan_result functionality of ECREnhancedScanVulnerabilityList class.
    It reads ecr_scan_result1.json which represents the vulnerabilities in ECR format and converts it to allowlist format using
    the construct_allowlist_from_ecr_scan_result method of ECREnhancedScanVulnerabilityList object. It then checks if the
    new vulnerability list formed in the Allowlist Format matches the allowlist1.json file.
    """
    minimum_sev_threshold = "HIGH"
    with open("./sanity/resources/ecr_scan_result1.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list.construct_allowlist_from_ecr_scan_result(scan_results)

    with open(f"./sanity/resources/allowlist1.json", "r") as f:
        expected_result = json.load(f)
    assert expected_result == get_object_after_serialization(
        ecr_image_vulnerability_list.vulnerability_list
    ), "Lists do not match!!"
    assert expected_result == get_object_after_serialization(
        ecr_image_vulnerability_list.get_sorted_vulnerability_list()
    ), "Lists do not match!!"


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration(
    "Check if ECREnhancedScanVulnerabilityList class conducts subtraction successfully"
)
def test_subtraction_operator():
    """
    This method tests the subtraction functionality of ECREnhancedScanVulnerabilityList class.
    It reads two ecr scan result files - ecr_scan_result1.json and ecr_scan_result2.json in 2 different
    ECREnhancedScanVulnerabilityList objects and then subtracts Allowlist Formatted vulnerability_list2 from
    vulnerability_list1 and vice versa. The subtraction results are then verified by comparing them with the
    allowlist_one_minus_two.json and allowlist_two_minus_one.json files respectively.
    """
    minimum_sev_threshold = "HIGH"
    with open("./sanity/resources/ecr_scan_result1.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list1 = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list1.construct_allowlist_from_ecr_scan_result(scan_results)

    with open("./sanity/resources/ecr_scan_result2.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list2 = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list2.construct_allowlist_from_ecr_scan_result(scan_results)

    vuln_list_one_minus_two = ecr_image_vulnerability_list1 - ecr_image_vulnerability_list2
    vuln_list_two_minus_one = ecr_image_vulnerability_list2 - ecr_image_vulnerability_list1

    with open("./sanity/resources/allowlist_one_minus_two.json", "r") as f:
        one_minus_two_saved_results = json.load(f)

    with open("./sanity/resources/allowlist_two_minus_one.json", "r") as f:
        two_minus_one_saved_results = json.load(f)

    assert (
        get_object_after_serialization(vuln_list_one_minus_two.vulnerability_list)
        == one_minus_two_saved_results
    ), f"{vuln_list_one_minus_two.vulnerability_list} \n does not match \n {one_minus_two_saved_results}"
    assert (
        get_object_after_serialization(vuln_list_two_minus_one.vulnerability_list)
        == two_minus_one_saved_results
    ), f"{vuln_list_two_minus_one.vulnerability_list} \n does not match \n {two_minus_one_saved_results}"
    assert (ecr_image_vulnerability_list1 - ecr_image_vulnerability_list1) == None, "Failed"


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration("Check if Allowlist files are read successfully")
def test_allowlist_file_read():
    minimum_sev_threshold = "HIGH"
    ecr_image_vulnerability_list = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list.construct_allowlist_from_file(
        "./sanity/resources/allowlist_to_read.json"
    )

    with open("./sanity/resources/dump_of_processed_allowlist_to_read.json", "r") as f:
        target_allowlist = json.load(f)

    assert (
        get_object_after_serialization(ecr_image_vulnerability_list.vulnerability_list)
        == target_allowlist
    ), f"{get_object_after_serialization(ecr_image_vulnerability_list.vulnerability_list)} \n does not match \n {target_allowlist} "


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration("Check if cmp operator works properly")
def test_compare_operator():
    """
    This method tests the compare i.e. __cmp__ functionality of ECREnhancedScanVulnerabilityList class.
    It reads ecr_scan_result1.json file in TWO different objects - object1a and object1b - and checks if
    __cmp__ returns True when applied on the 2 objects. It also reads ecr_scan_result2.json as object2 and
    then compares object2 with object1a and expects False as an outcome.
    """
    minimum_sev_threshold = "HIGH"

    with open("./sanity/resources/ecr_scan_result1.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list1a = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list1a.construct_allowlist_from_ecr_scan_result(scan_results)

    with open("./sanity/resources/ecr_scan_result1.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list1b = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list1b.construct_allowlist_from_ecr_scan_result(scan_results)

    with open("./sanity/resources/ecr_scan_result2.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list2 = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list2.construct_allowlist_from_ecr_scan_result(scan_results)

    assert (
        ecr_image_vulnerability_list1a == ecr_image_vulnerability_list1b
    ), "Same vulnerability lists not marked same"
    assert (
        ecr_image_vulnerability_list1a != ecr_image_vulnerability_list2
    ), "Different vulnerability lists marked same"


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration("Check if Allowlist files are read successfully")
def test_add_operator():
    """
    This method tests the addition functionality of ECREnhancedScanVulnerabilityList class.
    It reads two ecr scan result files - ecr_scan_result1.json and ecr_scan_result2.json in 2 different
    ECREnhancedScanVulnerabilityList objects and then adds the 2 objects to see if the results obtained
    is same as the one in allowlist_summation_result.json file.
    """
    minimum_sev_threshold = "HIGH"

    with open("./sanity/resources/ecr_scan_result1.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list1 = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list1.construct_allowlist_from_ecr_scan_result(scan_results)

    with open("./sanity/resources/ecr_scan_result2.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list2 = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    ecr_image_vulnerability_list2.construct_allowlist_from_ecr_scan_result(scan_results)

    ecr_image_vulnerability_list_one_plus_two = (
        ecr_image_vulnerability_list1 + ecr_image_vulnerability_list2
    )
    ecr_image_vulnerability_list_two_plus_one = (
        ecr_image_vulnerability_list2 + ecr_image_vulnerability_list1
    )

    with open("./sanity/resources/allowlist_summation_result.json", "r") as f:
        summation_allowlist = json.load(f)

    assert (
        get_object_after_serialization(
            ecr_image_vulnerability_list_one_plus_two.get_sorted_vulnerability_list()
        )
        == summation_allowlist
    ), f"Allowlist {ecr_image_vulnerability_list_one_plus_two.get_sorted_vulnerability_list()} does not match"
    assert (
        get_object_after_serialization(
            ecr_image_vulnerability_list_two_plus_one.get_sorted_vulnerability_list()
        )
        == summation_allowlist
    ), f"Allowlist {ecr_image_vulnerability_list_two_plus_one.get_sorted_vulnerability_list()} does not match"


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration("Compare Allowlist json list construction with ECR.")
def test_if_vulnerability_list_construction_from_ecr_results_and_allowlist_files_work_the_same():
    """
    This test ensures that the `vulnerability_list` constructed using ECR Enhanced Scans and the ones constructed using saved Allowlists
    have the same format and behave similarly. We read a saved allowlist (allowlist_to_read.json) and ecr_scan_result1.json and check if
    they are equal or not. In our case, we have ensured that both turn out to be equal - however, it is not always necessary. The importance
    of this test is to just ensure that the construct_allowlist_from_file leads to proper format of dataclass objects.
    """
    minimum_sev_threshold = "HIGH"

    with open("./sanity/resources/ecr_scan_result1.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list_constructed_from_ecr_scan_results = (
        ECREnhancedScanVulnerabilityList(minimum_severity=CVESeverity[minimum_sev_threshold])
    )
    ecr_image_vulnerability_list_constructed_from_ecr_scan_results.construct_allowlist_from_ecr_scan_result(
        scan_results
    )

    ecr_image_vulnerability_list_constructed_from_allowlist_json_file = (
        ECREnhancedScanVulnerabilityList(minimum_severity=CVESeverity[minimum_sev_threshold])
    )
    ecr_image_vulnerability_list_constructed_from_allowlist_json_file.construct_allowlist_from_file(
        "./sanity/resources/allowlist_to_read.json"
    )

    assert (
        ecr_image_vulnerability_list_constructed_from_ecr_scan_results
        == ecr_image_vulnerability_list_constructed_from_allowlist_json_file
    ), "Construcion of vulnerability_list from ecr scan results lead to different results as compared to construction of vulnerability_list from the allowlist JSON files."


@pytest.mark.usefixtures("sagemaker", "functionality_sanity")
@pytest.mark.model("N/A")
@pytest.mark.integration("FutureAllowlist")
def test_future_allowlist_generation():
    """
    In this test, we simulate the generation of the future allowlist.
    The latest image scan is stored in ecr_scan_result3.json and is read from there.
    The latest scan marks packages 2,7,6,8 and 9 as vulnerable.
    The already stored allowlist is github is stored in `allowlist_on_git_repo.json` and is read from there.
    The already stored allowlist shows packages 1,2,6 and 7 as allowlisted.
    The latest scan has the source_url parameter that is slightly different from the allowlist data, thus, we want package6's latest representation in future allowlist.
    The non_patchable vulnerabilities are store in non_patchable_vulns_data.json and consist of package6 and package9.
    The future allowlist is generated after reading all the above allowlists and then calling the generate_future_allowlist method.
    The genreated future allowlist is then compared with the `future_allowlist.json` file to make sure that packages 2,6,7 and 9 exist in future allowlist.
    This helps us test that the non-relevant vuln of Package 1 is removed from the allowlist.
    This helps us test that the package 6 in the future allowlist reflects the latest source_url and that in case of slight difference in ecr scan result and github allowlist data, the latest data from ecr scan is taken.
    This helps us test that the relevant vulns for Package 2 and 7 are preserved in the allowlist. Further, the reason_to_ignore for vulnerabilities comes from the manual allowlist if available.
    """
    minimum_sev_threshold = "CRITICAL" if is_huggingface_image() else "HIGH"

    with open("./sanity/resources/ecr_scan_result3.json", "r") as f:
        scan_results = json.load(f)
    ecr_image_vulnerability_list_constructed_from_ecr_scan_results = (
        ECREnhancedScanVulnerabilityList(minimum_severity=CVESeverity[minimum_sev_threshold])
    )
    ecr_image_vulnerability_list_constructed_from_ecr_scan_results.construct_allowlist_from_ecr_scan_result(
        scan_results
    )
    image_allowlist_created_from_github_file = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    image_allowlist_created_from_github_file.construct_allowlist_from_file(
        "./sanity/resources/allowlist_on_git_repo.json"
    )
    non_patchable_vulnerabilities = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    non_patchable_vulnerabilities.construct_allowlist_from_file(
        "./sanity/resources/non_patchable_vulns_data.json"
    )
    future_allowlist = generate_future_allowlist(
        ecr_image_vulnerability_list=ecr_image_vulnerability_list_constructed_from_ecr_scan_results,
        image_scan_allowlist=image_allowlist_created_from_github_file,
        non_patchable_vulnerabilities=non_patchable_vulnerabilities,
    )

    stored_future_allowlist_for_comparison = ECREnhancedScanVulnerabilityList(
        minimum_severity=CVESeverity[minimum_sev_threshold]
    )
    stored_future_allowlist_for_comparison.construct_allowlist_from_file(
        "./sanity/resources/future_allowlist.json"
    )
    assert (
        future_allowlist == stored_future_allowlist_for_comparison
    ), "Incorrect Future Allowlist generated"
